import tushare as ts
import pandas as pd
from typing import Dict, List, Optional
import numpy as np
from datetime import datetime, timedelta

class FundamentalData:
    """基本面数据类
    处理财务数据、估值数据和业绩预告等基本面相关数据
    """
    
    def __init__(self, ts_token: str):
        """初始化基本面数据类
        
        Args:
            ts_token: Tushare API token
        """
        self.ts_pro = ts.pro_api(ts_token)
    
    def get_financial_data(self, stock_code: str, period: str) -> Optional[Dict]:
        """获取财务数据
        
        Args:
            stock_code: 股票代码
            period: 报告期，如20200331表示2020年第一季度
            
        Returns:
            包含财务数据的字典
        """
        try:
            # 获取利润表数据
            income = self.ts_pro.income(
                ts_code=stock_code,
                period=period
            )
            
            # 获取资产负债表数据
            balance = self.ts_pro.balancesheet(
                ts_code=stock_code,
                period=period
            )
            
            # 获取现金流量表数据
            cashflow = self.ts_pro.cashflow(
                ts_code=stock_code,
                period=period
            )
            
            if income.empty and balance.empty and cashflow.empty:
                return None
                
            result = {
                'period': period,
                'stock_code': stock_code
            }
            
            # 提取关键财务指标
            if not income.empty:
                result.update({
                    'revenue': income['revenue'].values[0],  # 营业收入
                    'n_income': income['n_income'].values[0],  # 净利润
                    'total_profit': income['total_profit'].values[0],  # 利润总额
                    'oper_cost': income['oper_cost'].values[0]  # 营业成本
                })
            
            if not balance.empty:
                result.update({
                    'total_assets': balance['total_assets'].values[0],  # 总资产
                    'total_liab': balance['total_liab'].values[0],  # 总负债
                    'total_hldr_eqy_exc_min_int': balance['total_hldr_eqy_exc_min_int'].values[0]  # 股东权益
                })
            
            if not cashflow.empty:
                result.update({
                    'n_cashflow_act': cashflow['n_cashflow_act'].values[0],  # 经营活动现金流量净额
                    'n_cashflow_inv_act': cashflow['n_cashflow_inv_act'].values[0],  # 投资活动现金流量净额
                    'n_cash_flows_fnc_act': cashflow['n_cash_flows_fnc_act'].values[0]  # 筹资活动现金流量净额
                })
                
            return result
        except Exception as e:
            print(f"获取财务数据失败: {e}")
            return None
    
    def get_valuation_data(self, stock_code: str, trade_date: str) -> Optional[Dict]:
        """获取估值数据
        
        Args:
            stock_code: 股票代码
            trade_date: 交易日期，格式YYYYMMDD
            
        Returns:
            包含估值数据的字典
        """
        try:
            df = self.ts_pro.daily_basic(
                ts_code=stock_code,
                trade_date=trade_date,
                fields='ts_code,trade_date,pe,pe_ttm,pb,ps,ps_ttm,dv_ratio,dv_ttm,total_mv,circ_mv'
            )
            
            if df.empty:
                return None
                
            return {
                'trade_date': df['trade_date'].values[0],
                'pe': df['pe'].values[0],  # 市盈率
                'pe_ttm': df['pe_ttm'].values[0],  # 滚动市盈率
                'pb': df['pb'].values[0],  # 市净率
                'ps': df['ps'].values[0],  # 市销率
                'ps_ttm': df['ps_ttm'].values[0],  # 滚动市销率
                'dv_ratio': df['dv_ratio'].values[0],  # 股息率
                'dv_ttm': df['dv_ttm'].values[0],  # 滚动股息率
                'total_mv': df['total_mv'].values[0],  # 总市值
                'circ_mv': df['circ_mv'].values[0]  # 流通市值
            }
        except Exception as e:
            print(f"获取估值数据失败: {e}")
            return None
    
    def get_forecast(self, stock_code: str, period: str) -> Optional[Dict]:
        """获取业绩预告数据
        
        Args:
            stock_code: 股票代码
            period: 报告期，如20200331表示2020年第一季度
            
        Returns:
            包含业绩预告数据的字典
        """
        try:
            df = self.ts_pro.forecast(
                ts_code=stock_code,
                period=period
            )
            
            if df.empty:
                return None
                
            return {
                'period': df['period'].values[0],
                'ann_date': df['ann_date'].values[0],  # 公告日期
                'type': df['type'].values[0],  # 预告类型
                'p_change_min': df['p_change_min'].values[0],  # 预告净利润变动幅度下限
                'p_change_max': df['p_change_max'].values[0],  # 预告净利润变动幅度上限
                'net_profit_min': df['net_profit_min'].values[0],  # 预告净利润下限
                'net_profit_max': df['net_profit_max'].values[0],  # 预告净利润上限
                'summary': df['summary'].values[0]  # 业绩预告摘要
            }
        except Exception as e:
            print(f"获取业绩预告数据失败: {e}")
            return None
    
    def get_industry_pe(self, industry: str, trade_date: str) -> Optional[Dict]:
        """获取行业平均市盈率
        
        Args:
            industry: 行业名称
            trade_date: 交易日期，格式YYYYMMDD
            
        Returns:
            包含行业平均市盈率数据的字典
        """
        try:
            # 获取行业成分股
            stocks = self.ts_pro.stock_basic(
                exchange='',
                list_status='L',
                fields='ts_code,industry'
            )
            
            industry_stocks = stocks[stocks['industry'] == industry]['ts_code'].tolist()
            
            if not industry_stocks:
                return None
                
            # 获取成分股估值数据
            pe_list = []
            pb_list = []
            ps_list = []
            
            for stock in industry_stocks:
                try:
                    valuation = self.get_valuation_data(stock, trade_date)
                    if valuation and valuation['pe'] > 0:
                        pe_list.append(valuation['pe'])
                    if valuation and valuation['pb'] > 0:
                        pb_list.append(valuation['pb'])
                    if valuation and valuation['ps'] > 0:
                        ps_list.append(valuation['ps'])
                except:
                    continue
            
            if not pe_list:
                return None
                
            return {
                'industry': industry,
                'trade_date': trade_date,
                'avg_pe': np.mean(pe_list),
                'median_pe': np.median(pe_list),
                'avg_pb': np.mean(pb_list) if pb_list else None,
                'median_pb': np.median(pb_list) if pb_list else None,
                'avg_ps': np.mean(ps_list) if ps_list else None,
                'median_ps': np.median(ps_list) if ps_list else None
            }
        except Exception as e:
            print(f"获取行业平均市盈率失败: {e}")
            return None
            
    def get_profitability_indicators(self, stock_code: str, period: str) -> Optional[Dict]:
        """获取盈利能力指标
        
        Args:
            stock_code: 股票代码
            period: 报告期，如20200331表示2020年第一季度
            
        Returns:
            包含ROE、毛利率、净利润增长率等盈利能力指标的字典
        """
        try:
            # 获取当期财务数据
            current_data = self.get_financial_data(stock_code, period)
            if not current_data:
                return None
                
            # 计算上一年同期的报告期
            year = int(period[:4])
            quarter = period[4:]
            prev_year = year - 1
            prev_period = f"{prev_year}{quarter}"
            
            # 获取上一年同期财务数据（用于计算增长率）
            prev_data = self.get_financial_data(stock_code, prev_period)
            
            result = {
                'period': period,
                'stock_code': stock_code
            }
            
            # 计算ROE（净资产收益率）
            if 'n_income' in current_data and 'total_hldr_eqy_exc_min_int' in current_data and current_data['total_hldr_eqy_exc_min_int'] != 0:
                roe = current_data['n_income'] / current_data['total_hldr_eqy_exc_min_int'] * 100
                result['roe'] = roe
            
            # 计算毛利率
            if 'revenue' in current_data and 'oper_cost' in current_data and current_data['revenue'] != 0:
                gross_margin = (current_data['revenue'] - current_data['oper_cost']) / current_data['revenue'] * 100
                result['gross_margin'] = gross_margin
            
            # 计算净利润率
            if 'n_income' in current_data and 'revenue' in current_data and current_data['revenue'] != 0:
                net_profit_margin = current_data['n_income'] / current_data['revenue'] * 100
                result['net_profit_margin'] = net_profit_margin
            
            # 计算净利润增长率
            if prev_data and 'n_income' in current_data and 'n_income' in prev_data and prev_data['n_income'] != 0:
                net_profit_growth = (current_data['n_income'] - prev_data['n_income']) / abs(prev_data['n_income']) * 100
                result['net_profit_growth'] = net_profit_growth
            
            # 计算营收增长率
            if prev_data and 'revenue' in current_data and 'revenue' in prev_data and prev_data['revenue'] != 0:
                revenue_growth = (current_data['revenue'] - prev_data['revenue']) / prev_data['revenue'] * 100
                result['revenue_growth'] = revenue_growth
                
            return result
        except Exception as e:
            print(f"获取盈利能力指标失败: {e}")
            return None
    
    def get_solvency_indicators(self, stock_code: str, period: str) -> Optional[Dict]:
        """获取偿债能力指标
        
        Args:
            stock_code: 股票代码
            period: 报告期，如20200331表示2020年第一季度
            
        Returns:
            包含资产负债率、流动比率、速动比率等偿债能力指标的字典
        """
        try:
            # 获取资产负债表数据
            balance = self.ts_pro.balancesheet(
                ts_code=stock_code,
                period=period
            )
            
            if balance.empty:
                return None
                
            result = {
                'period': period,
                'stock_code': stock_code
            }
            
            # 计算资产负债率
            if 'total_assets' in balance.columns and 'total_liab' in balance.columns and balance['total_assets'].values[0] != 0:
                debt_to_assets = balance['total_liab'].values[0] / balance['total_assets'].values[0] * 100
                result['debt_to_assets'] = debt_to_assets
            
            # 计算流动比率
            if 'total_cur_assets' in balance.columns and 'total_cur_liab' in balance.columns and balance['total_cur_liab'].values[0] != 0:
                current_ratio = balance['total_cur_assets'].values[0] / balance['total_cur_liab'].values[0]
                result['current_ratio'] = current_ratio
            
            # 计算速动比率
            if 'total_cur_assets' in balance.columns and 'inventories' in balance.columns and 'total_cur_liab' in balance.columns and balance['total_cur_liab'].values[0] != 0:
                quick_assets = balance['total_cur_assets'].values[0] - balance['inventories'].values[0]
                quick_ratio = quick_assets / balance['total_cur_liab'].values[0]
                result['quick_ratio'] = quick_ratio
            
            return result
        except Exception as e:
            print(f"获取偿债能力指标失败: {e}")
            return None
    
    def get_free_cash_flow(self, stock_code: str, period: str) -> Optional[Dict]:
        """获取自由现金流数据
        
        Args:
            stock_code: 股票代码
            period: 报告期，如20200331表示2020年第一季度
            
        Returns:
            包含自由现金流数据的字典
        """
        try:
            # 获取现金流量表数据
            cashflow = self.ts_pro.cashflow(
                ts_code=stock_code,
                period=period
            )
            
            if cashflow.empty:
                return None
                
            # 获取资本支出数据
            capex = abs(cashflow['fix_intan_other_cash_pay'].values[0]) if 'fix_intan_other_cash_pay' in cashflow.columns else 0
            
            # 计算自由现金流 = 经营活动现金流 - 资本支出
            fcf = cashflow['n_cashflow_act'].values[0] - capex
            
            return {
                'period': period,
                'stock_code': stock_code,
                'operating_cash_flow': cashflow['n_cashflow_act'].values[0],
                'capex': capex,
                'free_cash_flow': fcf
            }
        except Exception as e:
            print(f"获取自由现金流数据失败: {e}")
            return None
    
    def get_peg(self, stock_code: str, trade_date: str, period: str) -> Optional[Dict]:
        """获取PEG指标数据
        
        Args:
            stock_code: 股票代码
            trade_date: 交易日期，格式YYYYMMDD
            period: 报告期，格式YYYYMMDD
            
        Returns:
            包含PEG指标数据的字典
        """
        try:
            # 获取PE数据
            valuation = self.get_valuation_data(stock_code, trade_date)
            if not valuation or not valuation['pe_ttm']:
                return None
            
            # 获取净利润增长率
            profitability = self.get_profitability_indicators(stock_code, period)
            if not profitability or 'net_profit_growth' not in profitability:
                return None
            
            # 计算PEG
            peg = valuation['pe_ttm'] / profitability['net_profit_growth'] if profitability['net_profit_growth'] != 0 else None
            
            return {
                'trade_date': trade_date,
                'stock_code': stock_code,
                'pe_ttm': valuation['pe_ttm'],
                'net_profit_growth': profitability['net_profit_growth'],
                'peg': peg
            }
        except Exception as e:
            print(f"获取PEG指标数据失败: {e}")
            return None
    
    def get_industry_prosperity(self, industry: str) -> Optional[Dict]:
        """获取行业景气度数据
        
        Args:
            industry: 行业名称
            
        Returns:
            包含行业景气度数据的字典
        """
        try:
            # 获取行业政策新闻
            news = self.ts_pro.news(
                src='新浪财经',
                start_date=(datetime.now() - timedelta(days=90)).strftime('%Y%m%d'),
                end_date=datetime.now().strftime('%Y%m%d')
            )
            
            # 筛选行业相关新闻
            industry_news = news[news['content'].str.contains(industry, na=False)]
            
            # 获取行业相关政策数量
            policy_keywords = ['政策', '规划', '补贴', '支持', '发展']
            policy_news = industry_news[industry_news['content'].str.contains('|'.join(policy_keywords), na=False)]
            
            # 获取行业集中度数据
            stocks = self.ts_pro.stock_basic(
                exchange='',
                list_status='L',
                fields='ts_code,industry'
            )
            industry_stocks = len(stocks[stocks['industry'] == industry])
            
            return {
                'industry': industry,
                'news_count': len(industry_news),
                'policy_news_count': len(policy_news),
                'industry_stocks_count': industry_stocks
            }
        except Exception as e:
            print(f"获取行业景气度数据失败: {e}")
            return None
    
    def get_company_moat(self, stock_code: str, period: str) -> Optional[Dict]:
        """获取公司护城河数据
        
        Args:
            stock_code: 股票代码
            period: 报告期，如20200331表示2020年第一季度
            
        Returns:
            包含公司护城河数据的字典
        """
        try:
            # 获取研发投入数据
            income = self.ts_pro.income(
                ts_code=stock_code,
                period=period
            )
            
            if income.empty:
                return None
            
            # 计算研发投入占比
            rd_expense = income['rd_exp'].values[0] if 'rd_exp' in income.columns else 0
            rd_ratio = rd_expense / income['revenue'].values[0] * 100 if income['revenue'].values[0] != 0 else 0
            
            # 获取专利数据（这里需要补充实际的专利数据获取逻辑）
            patents_count = 0  # 示例数据
            
            # 获取市场份额数据（这里需要补充实际的市场份额数据获取逻辑）
            market_share = 0  # 示例数据
            
            return {
                'period': period,
                'stock_code': stock_code,
                'rd_expense': rd_expense,
                'rd_ratio': rd_ratio,
                'patents_count': patents_count,
                'market_share': market_share
            }
        except Exception as e:
            print(f"获取公司护城河数据失败: {e}")
            return None